package com.alibaba.nacos.core.remote.grpc;

import io.grpc.*;
import io.grpc.internal.ServerStream;
import io.grpc.internal.ServerStreamHelper;
import io.grpc.netty.shaded.io.grpc.netty.NettyChannelHelper;
import io.grpc.netty.shaded.io.netty.channel.Channel;

import java.util.Objects;

public class GrpcConnectionInterceptor
        implements ServerInterceptor {
    @Override
    public <T, S> ServerCall.Listener<T> interceptCall(ServerCall<T, S> serverCall,
                                                       Metadata metadata,
                                                       ServerCallHandler<T, S> serverCallHandler) {
        Attributes attributes = serverCall.getAttributes();

        Context context = Context.current()
                .withValue(GrpcServerConstants.CONTEXT_KEY_CONN_ID, attributes.get(GrpcServerConstants.ATTR_TRANS_KEY_CONN_ID))
                .withValue(GrpcServerConstants.CONTEXT_KEY_CONN_REMOTE_IP, attributes.get(GrpcServerConstants.ATTR_TRANS_KEY_REMOTE_IP))
                .withValue(GrpcServerConstants.CONTEXT_KEY_CONN_REMOTE_PORT, attributes.get(GrpcServerConstants.ATTR_TRANS_KEY_REMOTE_PORT))
                .withValue(GrpcServerConstants.CONTEXT_KEY_CONN_LOCAL_PORT, attributes.get(GrpcServerConstants.ATTR_TRANS_KEY_LOCAL_PORT));

        if (Objects.equals(GrpcServerConstants.REQUEST_BI_STREAM_SERVICE_NAME, serverCall.getMethodDescriptor().getServiceName())) {
            context = context.withValue(GrpcServerConstants.CONTEXT_KEY_CHANNEL, getInternalChannel(serverCall));
        }

        return Contexts.interceptCall(context, serverCall, metadata, serverCallHandler);
    }

    private Channel getInternalChannel(ServerCall serverCall) {
        ServerStream serverStream = ServerStreamHelper.getServerStream(serverCall);
        Channel channel = NettyChannelHelper.getChannel(serverStream);
        return channel;
    }
}
